package com.artislong;

import cn.hutool.extra.servlet.JakartaServletUtil;
import cn.hutool.http.ContentType;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.server.transport.HttpServletSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import jakarta.servlet.ServletConfig;
import jakarta.servlet.ServletException;
import jakarta.servlet.annotation.WebServlet;
import jakarta.servlet.http.HttpServlet;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;

import java.io.IOException;
import java.net.http.HttpRequest;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

/**
 * @author 陈敏
 * @version McpClientServerServlet.java, v 1.0 2025 07 21 17:59 chenmin Exp $
 * Created on 2025 07 21
 */
@WebServlet(asyncSupported = true)
public class McpClientServerServlet extends HttpServlet {

    private McpSyncServer mcpSyncServer;
    private McpSyncClient mcpSyncClient;
    private HttpServletSseServerTransportProvider transportProvider;
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void init(ServletConfig config) throws ServletException {
        super.init(config);
        initMcpClient();
        initMcpServer();
    }

    private void initMcpClient() {

        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
                .header("Content-Type", "application/json")
                .timeout(Duration.ofMinutes(1));

        HttpClientSseClientTransport httpClientSseClientTransport = HttpClientSseClientTransport.builder("https://mcp.api-inference.modelscope.net")
                .sseEndpoint("/382d85b99d6947/sse")
                .objectMapper(objectMapper)
                .requestBuilder(requestBuilder)
                .build();

        McpSchema.ClientCapabilities clientCapabilities = McpSchema.ClientCapabilities.builder()
                .roots(true)      // Enable filesystem roots support with list changes notifications
                .sampling()       // Enable LLM sampling support
                .build();

        // Create a sync client with custom configuration
        mcpSyncClient = McpClient.sync(httpClientSseClientTransport)
                .clientInfo(new McpSchema.Implementation("my-client", "1.0.0"))
                .requestTimeout(Duration.ofMinutes(1))
                .initializationTimeout(Duration.ofMinutes(1))
                .capabilities(clientCapabilities)
                .build();

        // Initialize connection
        mcpSyncClient.initialize();
//        mcpSyncClient.setLoggingLevel(McpSchema.LoggingLevel.INFO);
    }

    private void initMcpServer() {
        McpSchema.ServerCapabilities serverCapabilities = McpSchema.ServerCapabilities.builder()
                .logging()
                .prompts(true)
                .resources(true, true)
                .tools(true)
                .completions()
                .experimental(new HashMap<>())
                .build();

        // 创建传输提供者
        transportProvider = HttpServletSseServerTransportProvider.builder()
                .objectMapper(new ObjectMapper())
                .baseUrl("/mcp")
                .messageEndpoint("/mcp/message")
                .sseEndpoint("/mcp/sse")
                .build();

        // 创建并配置MCP服务器
        mcpSyncServer = McpServer.sync(transportProvider)
                .serverInfo("my-server", "1.0.0")
                .capabilities(serverCapabilities)
                .requestTimeout(Duration.ofSeconds(10))
                .build();

        initTools();
    }

    private void initTools() {
        mcpSyncClient.listTools().tools().forEach(tool -> mcpSyncServer.addTool(toolSpecification(tool)));
    }

    private McpServerFeatures.SyncToolSpecification toolSpecification(McpSchema.Tool tool) {
        // 方法描述对应的执行逻辑
        BiFunction<McpSyncServerExchange, Map<String, Object>, McpSchema.CallToolResult> call = (exchange, arguments) -> mcpSyncClient.callTool(
                new McpSchema.CallToolRequest(tool.name(), arguments)
        );
        // 绑定描述和执行逻辑
        return new McpServerFeatures.SyncToolSpecification(tool, call);
    }

    @Override
    public void service(HttpServletRequest req, HttpServletResponse resp) throws ServletException, IOException {
        // 将请求委托给传输提供者处理
        transportProvider.service(req, resp);
    }

    public void get(HttpServletRequest req, HttpServletResponse resp) {
        resp.setCharacterEncoding("UTF-8");
        resp.setContentType("application/json");
        String query = req.getParameter("query");
        switch (query) {
            case "tools":
                McpSchema.ListToolsResult tools = mcpSyncClient.listTools();
                List<JSONObject> toolMetaDatas = tools.tools().stream().map(tool -> {
                    JSONObject result = new JSONObject();
                    result.set(tool.name(), tool.description());
                    result.set("schema", tool.inputSchema());
                    return result;
                }).collect(Collectors.toList());
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(toolMetaDatas), ContentType.JSON.getValue());
                break;
            case "prompts":
                McpSchema.ListPromptsResult prompts = mcpSyncClient.listPrompts();
                List<JSONObject> promptMetaDatas = prompts.prompts().stream().map(prompt -> {
                    JSONObject result = new JSONObject();
                    result.set(prompt.name(), prompt.description());
                    result.set("arguments", prompt.arguments());
                    return result;
                }).collect(Collectors.toList());
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(promptMetaDatas), ContentType.JSON.getValue());
                break;
            case "resources":
                McpSchema.ListResourcesResult resources = mcpSyncClient.listResources();
                List<JSONObject> resourceMetaDatas = resources.resources().stream().map(resource -> {
                    JSONObject result = new JSONObject();
                    result.set(resource.name(), resource.description());
                    result.set("mimeType", resource.mimeType());
                    result.set("uri", resource.uri());
                    result.set("annotations", resource.annotations());
                    return result;
                }).collect(Collectors.toList());
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(resourceMetaDatas), ContentType.JSON.getValue());
                break;
            case "call":
                String method = req.getParameter("method");
                Map<String, Object> paramMap = new HashMap<>(JakartaServletUtil.getParamMap(req));
                paramMap.remove("method");
                paramMap.remove("query");
                McpSchema.CallToolResult result = mcpSyncClient.callTool(
                        new McpSchema.CallToolRequest(method, paramMap)
                );
                List<McpSchema.Content> content = result.content();
                JakartaServletUtil.write(resp, JSONUtil.toJsonStr(content.get(0)), ContentType.JSON.getValue());
                break;
        }
    }

    @Override
    public void destroy() {
        // 关闭MCP服务器
        if (mcpSyncServer != null) {
            try {
                mcpSyncServer.close();
            } catch (Exception e) {
                getServletContext().log("Error closing MCP server", e);
            }
        }
        // 关闭MCP服务器
        if (mcpSyncClient != null) {
            try {
                mcpSyncClient.close();
            } catch (Exception e) {
                getServletContext().log("Error closing MCP server", e);
            }
        }
        super.destroy();
    }
}
